package org.ctstudio

import org.apache.commons.csv.CSVFormat
import org.apache.commons.csv.CSVParser
import org.apache.commons.csv.CSVPrinter
import org.apache.commons.math3.ml.clustering.CentroidCluster
import org.apache.commons.math3.ml.clustering.Clusterable
import org.apache.commons.math3.ml.clustering.KMeansPlusPlusClusterer
import org.ctstudio.extensions.assignEach
import org.ctstudio.extensions.minus
import java.io.FileReader
import java.io.FileWriter
import kotlin.math.max
import kotlin.math.min

class UserPV(var id: Int, private val pv: DoubleArray, var clusterId: Int = 0) : Clusterable {
    override fun getPoint(): DoubleArray {
        return pv
    }

    override fun toString(): String {
        return "{id:$id,point:${point.toList()}}"
    }
}

fun <T : Clusterable> normalize(points: List<T>, dimension: Int = points.first().point.size) {
    val maxAry = DoubleArray(dimension) { Double.NEGATIVE_INFINITY }
    val minAry = DoubleArray(dimension) { Double.POSITIVE_INFINITY }
    points.forEach {
        maxAry.assignEach { index, item -> max(item, it.point[index]) }
        minAry.assignEach { index, item -> min(item, it.point[index]) }
    }
    val denominator = maxAry - minAry
    points.forEach {
        // (x - min)/(max - min)
        it.point.assignEach { i, item -> (item - minAry[i]) / denominator[i] }
    }
}

fun main(args: Array<String>) {
    var inputFile = "visitor-statistics.csv"
    var categoryFile = "categories.csv"
    var fileCenters = "centers.csv"
    var fileClusters = "clusters.csv"
    if (args.size >= 4) {
        inputFile = args[0]
        categoryFile = args[1]
        fileCenters = args[2]
        fileClusters = args[3]
    }
    //1.读取数据
    val originData = loadData(inputFile)
    val filteredData = originData.filter { it.id > 0 }
    //2.异常数据处理
    filteredData.forEach { it.point.forEachIndexed { i, d -> if (d < 0.0) it.point[i] = 0.0 } }
    //3.数据归一化
    normalize(filteredData)
    val evaluator = CalinskiHarabasz<UserPV>()
    var maxScore = 0.0
    var bestKMeans: KMeansPlusPlusClusterer<UserPV>? = null
    var bestClusters: List<CentroidCluster<UserPV>>? = null
    //4.依次尝试k=2~10的聚类
    for (k in 2..10) {
        val kMeans = KMeansPlusPlusClusterer<UserPV>(k)
        val clusters = kMeans.cluster(filteredData)
        //5.评估并保留最佳聚类结果
        val score = evaluator.score(clusters)
        if (score > maxScore) {
            maxScore = score
            bestKMeans = kMeans
            bestClusters = clusters
        }
        println("k=$k,score=$score")
    }
    println("Best k is ${bestKMeans!!.k}")
    //6.1保存最佳聚类中心
    saveCenters(bestClusters!!, categoryFile, fileCenters)
    //6.2保存最佳聚类数据
    saveClusters(bestClusters, fileClusters)
}

fun loadData(filePath: String): List<UserPV> {
    val fmt = CSVFormat.EXCEL
    FileReader(filePath).use { reader ->
        return CSVParser.parse(reader, fmt).records.map {
            val uid = it.first().toIntOrNull() ?: 0
            val pv = DoubleArray(13) { i ->
                it[i + 1]?.toDoubleOrNull() ?: 0.0
            }
            UserPV(uid, pv)
        }
    }
}

fun writeCSV(filePath: String, mapper: (printer: CSVPrinter) -> Unit) {
    val fmt = CSVFormat.EXCEL
    FileWriter(filePath).use { writer ->
        CSVPrinter(writer, fmt).use { printer -> mapper(printer) }
    }
}

fun readCategories(fileCategories: String): List<String> {
    val fmt = CSVFormat.EXCEL
    FileReader(fileCategories).use { reader ->
        return CSVParser.parse(reader, fmt).records.first().toList()
    }
}

fun saveCenters(
    clusters: List<CentroidCluster<UserPV>>,
    fileCategories: String,
    fileCenters: String
) {
    val categories = readCategories(fileCategories)
    writeCSV(fileCenters) { printer ->
        printer.print("")
        printer.printRecord(categories)
        for (cluster in clusters) {
            //每类用户数
            printer.print(cluster.points.size)
            //每类访问量均值
            printer.printRecord(cluster.center.point.toList())
        }
    }
}

fun saveClusters(
    clusters: List<CentroidCluster<UserPV>>,
    fileClusters: String
) {
    writeCSV(fileClusters) { printer ->
        var clusterId = 0
        clusters.flatMap {
            clusterId++
            it.points.onEach { p -> p.clusterId = clusterId }
        }.sortedBy { it.id }.forEach { printer.printRecord(it.id, it.clusterId) }
    }
}